import subprocess
import argparse
import os

def build_perf_record_cmd(events=None, pids=None, tids=None, uid=None, system_wide=False,
                          cpus=None, rt_priority=None, count=None, freq=99,
                          output_file="perf.data", sleep_duration=30):
    # Base command
    cmd = ["perf", "record"]
    
    # Add options based on provided arguments
    if events:
        cmd.extend(["-e", events])
    if pids:
        cmd.extend(["-p", pids])
    if tids:
        cmd.extend(["-t", tids])
    if uid:
        cmd.extend(["-u", str(uid)])
    if system_wide:
        cmd.append("-a")
    cmd.append("-g")  # Enable call-graph by default
    if cpus:
        cmd.extend(["-C", cpus])
    if rt_priority is not None:
        cmd.extend(["-r", str(rt_priority)])
    if count is not None:
        cmd.extend(["-c", str(count)])
    if freq is not None:
        cmd.extend(["-F", str(freq)])
    if output_file:
        cmd.extend(["-o", output_file])
    
    # Add the command to run after the perf options
    cmd.append("--")
    cmd.append(f"sleep {sleep_duration}")

    return " ".join(cmd)

def generate_flamegraph(perf_record_cmd, output_file="perf.data"):
    base_name = os.path.splitext(output_file)[0]
    unfolded_file = f"{base_name}.unfold"
    folded_file = f"{base_name}.folded"
    svg_file = f"{base_name}.svg"

    # Step 1: Run perf record with the given command
    print("Running perf record...")
    subprocess.run(perf_record_cmd, shell=True, check=True)

    # Step 2: Generate unfolded stack traces
    print(f"Generating unfolded stack traces into {unfolded_file}...")
    perf_script_cmd = f"perf script -i {output_file} &> {unfolded_file}"
    subprocess.run(perf_script_cmd, shell=True, check=True)

    # Step 3: Collapse the stack traces
    print(f"Collapsing stack traces into {folded_file}...")
    stackcollapse_cmd = f"./FlameGraph/stackcollapse-perf.pl {unfolded_file} &> {folded_file}"
    subprocess.run(stackcollapse_cmd, shell=True, check=True)

    # Step 4: Generate the flame graph
    print(f"Generating flame graph into {svg_file}...")
    flamegraph_cmd = f"./FlameGraph/flamegraph.pl {folded_file} > {svg_file}"
    subprocess.run(flamegraph_cmd, shell=True, check=True)

    print(f"Flame graph saved as {svg_file}")

def main():
    parser = argparse.ArgumentParser(description="Generate a flame graph using perf and FlameGraph tools.")
    parser.add_argument("-e", "--events", help="Comma-separated list of events to record.")
    parser.add_argument("-p", "--pids", help="Comma-separated list of process IDs to profile.")
    parser.add_argument("-t", "--tids", help="Comma-separated list of thread IDs to profile.")
    parser.add_argument("-u", "--uid", help="User ID or name to profile.")
    parser.add_argument("-a", "--system-wide", action="store_true", help="Collect data from all CPUs.")
    parser.add_argument("-C", "--cpus", help="Comma-separated list of CPUs to profile.")
    parser.add_argument("-r", "--rt-priority", type=int, help="Real-time priority to run perf.")
    parser.add_argument("-c", "--count", type=int, help="Event count per sample.")
    parser.add_argument("-F", "--freq", type=int, default=99, help="Sampling frequency.")
    parser.add_argument("-o", "--output-file", default="perf.data", help="Output file name for perf data.")
    parser.add_argument("-s", "--sleep-duration", type=int, default=30, help="Duration for sleep command.")

    args = parser.parse_args()
    perf_record_cmd = build_perf_record_cmd(events=args.events, pids=args.pids, tids=args.tids, uid=args.uid,
                                            system_wide=args.system_wide, cpus=args.cpus, rt_priority=args.rt_priority,
                                            count=args.count, freq=args.freq, output_file=args.output_file,
                                            sleep_duration=args.sleep_duration)
    print(f"Generated perf record command: {perf_record_cmd}")
    generate_flamegraph(perf_record_cmd, output_file=args.output_file)

if __name__ == "__main__":
    main()
